import numpy as np
import PIL.Image as Image
import patchreconst as pr

PATCH_SIZE = 32
COMPONENT_COUNT = 100
PATCH_COUNT = 400000
EPOCH_COUNT = 16
CODE_COUNT = 800
SPARSITY_PARAMETER = 4.0
NON_NEGATIVE = True

fileNamesFile = open('files.txt')
fileNames = fileNamesFile.read().strip().split("\n")

scBasisFile = 'basis/basisnn' + str(NON_NEGATIVE) + 'lamb' + str(SPARSITY_PARAMETER) + 'comps' + str(COMPONENT_COUNT) + 'codes' + str(CODE_COUNT) + 'patches' + str(PATCH_COUNT) + 'epochs' + str(EPOCH_COUNT) + '.npy'
scBasis = np.load(scBasisFile)

icaFilterFile = 'icabasis/basisluisicacomps' + str(COMPONENT_COUNT) + 'codes' + str(CODE_COUNT) + 'patches' + str(PATCH_COUNT) + '.npy'
icaFilters = np.load(icaFilterFile)

scMSE = 0.0
icaMSE = 0.0
v1MSE = 0.0

for fileName in fileNames:
	image = np.load(fileName)
	imageWithNoise = np.load(fileName[:-4] + '_noise.npy')

	rowPad = 0

	if image.shape[0] % PATCH_SIZE != 0:
		lastPatchRows = image.shape[0] - (int(image.shape[0] / PATCH_SIZE) * PATCH_SIZE)
		rowPad = PATCH_SIZE - lastPatchRows

	if image.shape[1] % PATCH_SIZE != 0:
		lastPatchCols = image.shape[1] - (int(image.shape[1] / PATCH_SIZE) * PATCH_SIZE)
		colPad = PATCH_SIZE - lastPatchCols

	paddedImage = np.zeros((image.shape[0] + rowPad, image.shape[1] + colPad))
	paddedImage[:image.shape[0], :image.shape[1]] = image[:]

	paddedImage -= np.mean(paddedImage)
	paddedImage /= np.std(paddedImage)

	reconstV1 = np.zeros(paddedImage.shape)
	reconstSC = np.zeros(paddedImage.shape)
	reconstICA = np.zeros(paddedImage.shape)

	rows = int(paddedImage.shape[0] / PATCH_SIZE)
	cols = int(paddedImage.shape[1] / PATCH_SIZE)

	patches = np.zeros((rows, cols, PATCH_SIZE, PATCH_SIZE))

	for i in range(rows):
		for j in range(cols):
			row = i * PATCH_SIZE
			col = j * PATCH_SIZE
			rowEnd = row + PATCH_SIZE
			colEnd = col + PATCH_SIZE

			patches[i, j, :, :] = paddedImage[row:rowEnd, col:colEnd]

	patches = patches.reshape((-1, PATCH_SIZE, PATCH_SIZE))
	patches = patches.reshape((patches.shape[0], PATCH_SIZE * PATCH_SIZE))

	patches -= np.mean(patches, axis = -1)[:, np.newaxis]

	stds = np.std(patches, axis = -1)[:, np.newaxis]
	stds[np.where(stds == 0.0)] = 1.0

	patches /= stds

	patches = patches.reshape((-1, PATCH_SIZE, PATCH_SIZE))

	scCodes, icaCodes, v1Simple, v1cMean, angles = pr.responses(patches, scBasis, SPARSITY_PARAMETER, NON_NEGATIVE, icaFilters)

	patchesReconstV1 = pr.reconstructV1(v1Simple).reshape((rows, cols, PATCH_SIZE, PATCH_SIZE))
	patchesReconstSC = pr.reconstruct(scBasis, scCodes, v1cMean, angles).reshape((rows, cols, PATCH_SIZE, PATCH_SIZE))
	patchesReconstICA = pr.reconstruct(icaFilters, icaCodes, v1cMean, angles).reshape((rows, cols, PATCH_SIZE, PATCH_SIZE))

	for i in range(rows):
		for j in range(cols):
			row = i * PATCH_SIZE
			col = j * PATCH_SIZE
			rowEnd = row + PATCH_SIZE
			colEnd = col + PATCH_SIZE

		reconstV1[row:rowEnd, col:colEnd] = patchesReconstV1[i, j]
		reconstSC[row:rowEnd, col:colEnd] = patchesReconstSC[i, j]
		reconstICA[row:rowEnd, col:colEnd] = patchesReconstICA[i, j]

	reconstV1 = reconstV1[:image.shape[0], :image.shape[1]]
	reconstSC = reconstSC[:image.shape[0], :image.shape[1]]
	reconstICA = reconstICA[:image.shape[0], :image.shape[1]]

	pixelCount = float(image.shape[0] * image.shape[1])

	v1MSE += np.sum((reconstV1 - image) ** 2.0) / pixelCount
	scMSE += np.sum((reconstSC - image) ** 2.0) / pixelCount
	icaMSE += np.sum((reconstICA - image) ** 2.0) / pixelCount

fileCount = float(len(fileNames))

v1MSE /= fileCount
scMSE /= fileCount
icaMSE /= fileCount
